{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1d0a45",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.core.helper.session_helper import Session, get_execution_role\n",
    "\n",
    "sess = Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23466452",
   "metadata": {},
   "source": [
    "### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ff2dd50",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "train_dataset, test_dataset = load_dataset(\"imdb\", split=[\"train\", \"test\"])\n",
    "train_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20d80aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset[10]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fb4b06e",
   "metadata": {},
   "source": [
    "### Use FrameworkProcessor with pytorch image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba7160e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.core.image_uris import get_training_image_uri\n",
    "from sagemaker.core.processing import FrameworkProcessor\n",
    "\n",
    "image_uri = get_training_image_uri(\n",
    "    region=sess.boto_region_name,\n",
    "    framework=\"pytorch\",\n",
    "    framework_version=\"1.13\",\n",
    "    py_version=\"py39\",\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    ")\n",
    "\n",
    "pytorch_processor = FrameworkProcessor(\n",
    "    image_uri=image_uri,\n",
    "    role=role,\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    "    instance_count=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91093062",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.core.shapes import ProcessingOutput, ProcessingS3Output\n",
    "from time import gmtime, strftime\n",
    "import os\n",
    "\n",
    "s3_prefix = \"huggingface-text-classification\"\n",
    "processing_job_name = \"{}-{}\".format(s3_prefix, strftime(\"%d-%H-%M-%S\", gmtime()))\n",
    "output_destination = \"s3://{}/{}\".format(sess.default_bucket(), s3_prefix)\n",
    "\n",
    "pytorch_processor.run(\n",
    "    code=\"preprocessing.py\",\n",
    "    source_dir=os.path.abspath(\"scripts/preprocess\"),\n",
    "    job_name=processing_job_name,\n",
    "    outputs=[\n",
    "        ProcessingOutput(\n",
    "            output_name=\"train\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=\"{}/train\".format(output_destination),\n",
    "                local_path=\"/opt/ml/processing/train\",\n",
    "                s3_upload_mode=\"EndOfJob\",\n",
    "            ),\n",
    "        ),\n",
    "        ProcessingOutput(\n",
    "            output_name=\"test\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=\"{}/test\".format(output_destination),\n",
    "                local_path=\"/opt/ml/processing/test\",\n",
    "                s3_upload_mode=\"EndOfJob\",\n",
    "            ),\n",
    "        ),\n",
    "    ],\n",
    "    wait=False,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de7dde5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_processor.latest_job.refresh().processing_job_status"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
